import numpy as np
import matplotlib.pyplot as plt

# 西瓜数据集3.0alpha
# 密度 density
d = np.array([0.697, 0.774, 0.634, 0.608, 0.556, 0.403, 0.481, 0.437,
              0.666, 0.243, 0.245, 0.343, 0.639, 0.657, 0.360, 0.593, 0.719])
# 含糖率 sugar content
s = np.array([0.460, 0.376, 0.264, 0.318, 0.215, 0.237, 0.149, 0.211,
              0.091, 0.267, 0.057, 0.099, 0.161, 0.198, 0.370, 0.042, 0.103])
# 1 表示好瓜 0 表示坏瓜
label = np.array([1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0])

mean1 = np.array([np.mean(d[0:8]), np.mean(s[0:8])]) # 第1类（好瓜）均值向量
mean0 = np.array([np.mean(d[8:17]), np.mean(s[8:17])]) # 第0类（坏瓜）均值向量

# 计算类内散度矩阵（公式3.33）
Sw = np.zeros((2, 2))
for i in range(17):
    if i < 8:
        temp = np.mat(np.array([d[i], s[i]]) - mean1)
    else:
        temp = np.mat(np.array([d[i], s[i]]) - mean0)
    Sw += np.dot(temp.T, temp)

# 奇异值分解
U, Sigma, V = np.linalg.svd(Sw)
Sigma = np.diag(Sigma)
inv_Sw = np.dot(np.dot(V.T, np.linalg.inv(Sigma)), U.T)

# 权值计算（公式3.39）
w = np.dot(inv_Sw, mean0 - mean1)

# 测试
predict = [] # 线性判别分析预测结果
acc = 0 # 预测正确的数量
u1 = w[0] * mean1[0] + w[1] * mean1[1] # 好瓜的均值向量在直线上的投影点
u0 = w[0] * mean0[0] + w[1] * mean0[1] # 坏瓜的均值向量在直线上的投影点
for i in range(17):
    x = w[0] * d[i] + w[1] * s[i] # 第i个样本在直线上的投影点

    if (abs(x - u1) <= abs(x - u0) and label[i] == 1) or (abs(x - u1) > abs(x - u0) and label[i] == 0):
        # 预测正确
        acc += 1

    predict.append(1 if abs(x - u1) <= abs(x - u0) else 0)

print('Accuracy : ', float(acc) / float(len(label))) # 输出正确率
print('Label    : ', list(label)) # 输出标签
print('Predict  : ', predict) # 输出预测结果

# 绘图
x = np.arange(0, 0.9, 0.1)
# 投影直线为：0 = w0 * d + w1 * s
# 移项：s = (-w0/w1) * d
y = np.array((-w[0]/w[1]) * x)
good = plt.scatter(d[0:8], s[0:8], color = 'g', marker = 'o', label = 'good')
bad = plt.scatter(d[8:17], s[8:17], color = 'r', marker = 'x', label = 'bad')
plt.plot(x, y)
plt.xlabel('density')
plt.ylabel('sugar content')
plt.legend()
plt.savefig('./3_5.png')
